import json
import torch
from PIL import Image
from torch import nn
from torchvision import models, datasets, transforms
from torch.autograd import Variable
import os
from efficientnet_pytorch import EfficientNet
import matplotlib.pyplot as plt
from django.db.models import F
from rest_framework import serializers

import re
import base64
from io import BytesIO

from backend.api.models import FishInfo, FishInfoSerializer

transform = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])


class MyEncoder(json.JSONEncoder):

    def default(self, obj):
        """
        只要检查到了是bytes类型的数据就把它转为str类型
        :param obj:
        :return:
        """
        if isinstance(obj, bytes):
            return str(obj, encoding='utf-8')
        return json.JSONEncoder.default(self, obj)


class classifier(object):
    # 定义构造方法
    def __init__(self):  # __init__() 是类的初始化方法；它在类的实例化操作后 会自动调用，不需要手动调用；
        self.model = models.resnet101(pretrained=True)
        for parma in self.model.parameters():
            parma.requires_grad = False
        self.model.fc = torch.nn.Linear(2048, 10)
        resume = os.path.dirname(__file__) + '/model_best.path.tar'
        checkpoint = torch.load(resume)
        self.model.load_state_dict(checkpoint['state_dict'])
        model = self.model.cuda()
        model.eval()
        torch.no_grad()

    def predict(self, data):
        """鱼类识别接口"""
        # 解析数据
        base64_data = re.sub('^data:image/.+;base64,', '', data)
        num = len(base64_data) % 4
        if num != 0:
            base64_data = base64_data + '=' * (4 - num)
        byte_data = base64.b64decode(base64_data)
        image_data = BytesIO(byte_data)
        img = Image.open(image_data)
        img = transform(img).unsqueeze(0)
        img_ = Variable(img.cuda())
        # 调用模型做分类识别
        outputs = self.model(img_)
        _, indices = torch.max(outputs, 1)
        percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
        labels_map = json.load(open(os.path.dirname(__file__) + '/labels_fish_identification.txt'))
        labels_map = [labels_map[str(i)] for i in range(10)]
        _, indices = torch.sort(outputs, descending=True)  # obtain predicted result and sort form big to small
        percentage1 = torch.nn.functional.softmax(outputs, dim=1)[0] * 100  # reture percentage of predicted value
        result = [(labels_map[idx], percentage1[idx].item()) for idx in indices[0][:1]]
        print(result[0])
        # 从数据库查询鱼类信息将结果返回
        return self.get_fish_info(result[0][0]), result[0][1]

    def get_fish_info(self, label):
        """从数据库查询匹配的鱼类信息"""
        fish = FishInfo.objects.filter(label=label).first()
        response = FishInfoSerializer(instance=fish).data
        return response

    # 定义构造方法-本机测试
    # def __init__(self):
    #     self.model = models.resnet101(pretrained=True)
    #
    # def predict(self, data):
    #     return self.get_fish_info('a'), 0.4
